package com.hulk.rpc.utils.concurrent;

import java.util.Arrays;
import java.util.function.Supplier;

import static com.hulk.rpc.utils.TableUtils.offset;
import static com.hulk.rpc.utils.TableUtils.tableSizeFor;
import static com.hulk.rpc.utils.UnsafeUtils.unsafe;

/**
 * 高性能，仅适用于key值比较少的情况，不能超过512k个
 */
public class ConcurrentIntToObjectArrayMap<T> {

  private static final Object NOT_FOUND = new Object();
  private static final int MAXIMUM_CAPACITY = 1024 * 512;

  private static final int ABASE;
  private static final int ASHIFT;

  private volatile Object[] array;

  public ConcurrentIntToObjectArrayMap() {
    this(16);
  }

  public ConcurrentIntToObjectArrayMap(int initialCapacity) {
    if (initialCapacity < 2) {
      throw new IllegalArgumentException("Illegal initial capacity: " + initialCapacity);
    }

    if (initialCapacity > MAXIMUM_CAPACITY) {
      throw new IndexOutOfBoundsException("Illegal initial capacity: " + initialCapacity);
    }

    ensureCapacity(initialCapacity);
  }

  /**
   *
   */
  public boolean contains(int key) {
    if (key < 0) {
      return false;
    }

    Object[] finalArray = array;
    if (key >= finalArray.length) {
      return false;
    }

    Object value = unsafe().getObjectVolatile(finalArray, offset(ABASE, ASHIFT, key));

    return value != NOT_FOUND;
  }

  /**
   * @return null为找不到
   */
  @SuppressWarnings("unchecked")
  public T get(int key) {
    if (key < 0) {
      throw new IllegalArgumentException("Illegal key: " + key);
    }

    Object[] finalArray = array;
    if (key >= finalArray.length) {
      return null;
    }

    Object value = unsafe().getObjectVolatile(finalArray, offset(ABASE, ASHIFT, key));

    if (value != NOT_FOUND) {
      return (T) value;
    } else {
      return null;
    }
  }

  @SuppressWarnings("unchecked")
  public T getOrUpdate(int key, Supplier<T> producer) {
    if (key < 0) {
      throw new IllegalArgumentException("Illegal key: " + key);
    }

    Object value;
    final Object[] finalArray = array;

    if (key >= finalArray.length) {
      value = NOT_FOUND;
    } else {
      value = unsafe().getObjectVolatile(finalArray, offset(ABASE, ASHIFT, key));
    }

    if (value != NOT_FOUND) {
      return (T) value;
    }

    synchronized (this) {
      final Object[] theArray = array;
      if (key < theArray.length) {
        value = unsafe().getObjectVolatile(theArray, offset(ABASE, ASHIFT, key));
      }

      if (value != NOT_FOUND) {
        return (T) value;
      }

      value = producer.get();
      put(key, (T) value);

      return (T) value;
    }
  }

  /**
   * @param key 大于零，小于512k
   */
  public void put(int key, T value) {
    if (key < 0) {
      throw new IllegalArgumentException("Illegal key: " + key);
    }

    if (key >= MAXIMUM_CAPACITY) {
      throw new IndexOutOfBoundsException("Illegal key: " + key);
    }

    ensureCapacity(key + 1);
    final long offset = offset(ABASE, ASHIFT, key);

    for (; ; ) {// like cas
      final Object[] before = array;
      unsafe().putOrderedObject(before, offset, value);
      final Object[] after = array;

      if (before == after) {
        return;
      }
    }
  }

  /**
   * 重置为未赋值
   *
   * @param key 大于零，小于512k
   */
  public boolean remove(final int key) {
    if (key < 0) {
      return false;
    }

    Object[] finalArray = array;
    if (key >= finalArray.length) {
      return false;
    }

    final long offset = offset(ABASE, ASHIFT, key);

    for (; ; ) {// like cas
      final Object[] before = array;
      unsafe().putOrderedObject(before, offset, NOT_FOUND);
      final Object[] after = array;

      if (before == after) {
        return true;
      }
    }
  }

  public void clear() {
    if (array == null) {
      return;
    }

    Object[] objs = new Object[16];
    Arrays.fill(objs, NOT_FOUND);

    array = objs;
  }

  private void ensureCapacity(int capacity) {
    Object[] theArray = array;
    if (theArray != null && theArray.length >= capacity) {
      return;
    }

    synchronized (this) {
      Object[] finalArray = array;
      if (finalArray != null && finalArray.length >= capacity) {
        return;
      }

      int newCapacity = tableSizeFor(capacity);

      if (newCapacity > MAXIMUM_CAPACITY) {
        throw new IndexOutOfBoundsException(newCapacity);
      }

      Object[] objs = new Object[newCapacity];
      Arrays.fill(objs, NOT_FOUND);

      if (finalArray != null) {
        System.arraycopy(finalArray, 0, objs, 0, finalArray.length);
      }

      array = objs;
    }
  }

  static {
    try {
      ABASE = unsafe().arrayBaseOffset(Object[].class);

      int scale = unsafe().arrayIndexScale(Object[].class);
      if ((scale & (scale - 1)) != 0) {
        throw new Error("array index scale not a power of two");
      }

      ASHIFT = 31 - Integer.numberOfLeadingZeros(scale);
    } catch (Exception e) {
      throw new Error(e);
    }
  }

}
